(module inferrer  (lib "eopl.ss" "eopl")

  (require "drscheme-init.scm")
  (require "lang.scm")
  (require "static-data-structures.scm")
  (require "expand-type.scm")
  (require "unifier.scm")

  (provide type-of expand-answer)

  ;;;;;;;;;;;;;;;; The Type Checker ;;;;;;;;;;;;;;;;

  ;; we'll be thinking of the type of an expression as pair consisting
  ;; of a type (possibly with some type variables in it) and a
  ;; substitution that tells us how to interpret those type variables.

  ;; answer = type * subst
  ;; type-of: exp * tenv * subst * -> answer

  ;; we may change this once we get it working.

  (define-datatype answer answer?
    (return                        ; cunningly chosen constructor name
      (type type?)
      (subst substitution?)))

  (define expand-answer
    (lambda (ans)
      (cases answer ans
        (return (ty subst)
          (apply-subst subst ty)))))

  (define type-of
    (lambda (exp tenv subst)
      (cases expression exp

        ;; only change for module system
        (qualified-var-ref (m-name comp-name)
          (let ((m-type 
                  (expand-module-type-in-tenv
                    (lookup-module-name-in-tenv tenv m-name)
                    tenv)))
            (cases module-type m-type
              (env-module-type (decls)
                (return
                  (wrap-named-types-with-m-name
                    (lookup-val-decl-in-decls comp-name decls) 
                    m-name)
                  subst))
              (proc-module-type (name params result-type)
                (eopl:error 'type-of
                  "Can't extract variable from parameterized module ~s"
                  exp))
;;               (named-module-type (id)
;;                 (eopl:error 'type-of 
;;                   "expand-module-type-in-tenv returned a named interface ~s"
;;                   id))
                )))
            
        (var-exp (var)
          (return
            (lookup-variable-name-in-tenv tenv var)
            subst))

        (const-exp (num) (return (int-type) subst))
      
        (zero?-exp (exp1)
          (cases answer (type-of exp1 tenv subst)
            (return (type1 subst1)
              (let ((subst2 (unifier type1 (int-type) tenv subst exp)))
                (return (bool-type) subst2)))))

        (diff-exp (exp1 exp2)
          (cases answer (type-of exp1 tenv subst)
            (return (type1 subst1)
              (let ((subst1 (unifier type1 (int-type) tenv subst1 exp1)))
                (cases answer (type-of exp2 tenv subst1)
                  (return (type2 subst2)
                    (let ((subst2
                            (unifier type2 (int-type) tenv subst2 exp2)))
                      (return (int-type) subst2))))))))

        (if-exp (test-exp true-exp false-exp)
          (cases answer (type-of test-exp tenv subst)
            (return (test-type subst)
              (let ((subst (unifier test-type (bool-type) tenv subst
                             test-exp)))
                (cases answer (type-of true-exp tenv subst)
                  (return (true-type subst)
                    (cases answer (type-of false-exp tenv
                                    subst)
                      (return (false-type subst)
                        (let ((subst (unifier true-type false-type
                                       tenv subst exp)))
                          (return true-type subst))))))))))

        (let-exp (var exp1 body)
          (cases answer (type-of exp1 tenv subst)
            (return (rhs-type subst)
              (type-of body
                (extend-tenv var rhs-type tenv)
                subst))))

        (proc-exp (bvar otype body)
;;          (eopl:printf "inferring type for ~s~%" exp)
          (let ((arg-type (expand-type-in-tenv (otype->type otype) tenv)))
;;            (eopl:printf "expanded arg type is ~s~%" arg-type)
            (let ((body-type
                    (type-of body
                      (extend-tenv bvar arg-type tenv)
                      subst)))
              (cases answer body-type
                (return (result-type subst)
;;                  (eopl:printf "expanded body type is ~s~%"
;;                    result-type)
                  (let ((ty (proc-type 
                              arg-type
                              result-type
                              ;(expand-type-in-tenv result-type tenv)
                              )))
;;                    (eopl:printf "expression type is ~s~%" ty)
                  (return ty subst)))))))

        (call-exp (rator rand)
          (let ((result-type (fresh-tvar-type)))
            (cases answer (type-of rator tenv subst)
              (return (rator-type subst)
                (cases answer (type-of rand tenv subst)
                  (return (rand-type subst)
                    (let ((subst
                            (unifier rator-type
                              (proc-type rand-type result-type)
                              tenv
                              subst
                              exp)))
                      (return result-type subst))))))))

        (letrec-exp (proc-result-otype proc-name 
                      bvar proc-arg-otype 
                      proc-body
                      letrec-body)
          (let ((proc-result-type
                  (otype->type proc-result-otype)) 
                (proc-arg-type
                  (otype->type proc-arg-otype)))
            (let ((tenv-for-letrec-body
                    (extend-tenv
                      proc-name
                      (proc-type proc-arg-type proc-result-type)
                      tenv)))
              (cases answer (type-of proc-body
                              (extend-tenv
                                bvar proc-arg-type tenv-for-letrec-body)
                              subst)
                (return (proc-body-type subst)
                  (let ((subst 
                          (unifier proc-body-type proc-result-type
                            tenv subst
                            proc-body))) 
                    (type-of letrec-body
                      tenv-for-letrec-body
                      subst)))))))
        
        )))

;  (trace type-of)

  (define fresh-tvar-type
    (let ((sn 0))
      (lambda ()
        (set! sn (+ sn 1))
        (tvar-type sn))))

  (define otype->type
    (lambda (otype)
      (cases optional-type otype
        (no-type () (fresh-tvar-type))
        (a-type (ty) ty))))

  (define lookup-val-decl-in-decls
    (lambda (search-sym decls)
      (let ((maybe-answer
              (variable-name->maybe-binding-in-decls decls search-sym)))
        (if maybe-answer maybe-answer
          (eopl:error 'lookup-variable-name-in-decls
            "Unbound variable name ~s in decls ~s"
            search-sym decls)))))

  ;; these are decls, so innermost binding is on the right.
  (define variable-name->maybe-binding-in-decls
    (lambda (decls search-sym)
      (let recur ((decls decls))
        (if (null? decls) #f
          (let ((decl (car decls))
                (other-decls (cdr decls)))
            (cases declaration decl
              (val-declaration (name ty)
                (cond
                  ;; is search-sym bound to the right?
                  ((recur other-decls) => (lambda (x) x))
                  ;; is this it?
                  ((eqv? name search-sym) ty)
                  (else #f)))
              (abstract-type-declaration (name) (recur other-decls))
              (type-abbrev-declaration (name ty) (recur other-decls))
              ))))))

  )
